import argparse
import threading

import numpy as np

import genesis as gs


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-v", "--vis", action="store_true", default=False)
    args = parser.parse_args()

    ########################## init ##########################
    gs.init(backend=gs.cpu)

    ########################## create a scene ##########################
    viewer_options = gs.options.ViewerOptions(
        camera_pos=(2.5, 0.0, 1.5),
        camera_lookat=(0.0, 0.0, 0.5),
        camera_fov=30,
        max_FPS=60,
    )

    scene = gs.Scene(
        sim_options=gs.options.SimOptions(
            dt=0.01,
        ),
        viewer_options=viewer_options,
        show_viewer=args.vis,
    )

    ########################## entities ##########################
    plane = scene.add_entity(
        gs.morphs.Plane(),
    )
    drone = scene.add_entity(
        morph=gs.morphs.Drone(
            file="urdf/drones/cf2x.urdf",
            pos=(0.0, 0, 0.02),
        ),
    )

    ########################## build ##########################
    scene.build()

    traj = np.array(
        [
            [1.0, 1.0, 0.98824805, 1.0],
            [0.67815816, 1.0, 1.0, 1.0],
            [1.0, 0.87905186, 0.8319297, 1.0],
            [1.0, 0.85295373, 0.94554883, 1.0],
            [1.0, 1.0, 0.9663153, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 0.9954323, 1.0, 1.0],
            [1.0, 0.9974212, 1.0, 1.0],
            [1.0, 0.99529535, 1.0, 1.0],
            [1.0, 0.9965133, 1.0, 1.0],
            [1.0, 0.99495167, 1.0, 1.0],
            [1.0, 0.99533206, 1.0, 1.0],
            [1.0, 0.9941533, 1.0, 1.0],
            [1.0, 0.9937679, 1.0, 1.0],
            [1.0, 0.9926078, 1.0, 1.0],
            [1.0, 0.99150425, 1.0, 1.0],
            [1.0, 0.9899133, 1.0, 1.0],
            [1.0, 0.9879518, 1.0, 1.0],
            [1.0, 0.985401, 1.0, 0.93023926],
            [1.0, 0.9757207, 1.0, 0.8165948],
            [1.0, 0.94503635, 1.0, 0.47085124],
            [1.0, 0.9973584, 1.0, 0.92234856],
            [1.0, 0.97841257, 1.0, 0.9260282],
            [1.0, 0.9749091, 1.0, 0.8766382],
            [1.0, 0.9661152, 1.0, 0.8508391],
            [1.0, 0.9588664, 1.0, 0.8184431],
            [1.0, 0.9513355, 1.0, 0.78656846],
            [1.0, 0.9438352, 1.0, 0.7555151],
            [1.0, 0.9364986, 1.0, 0.7264303],
            [1.0, 0.92944163, 1.0, 0.70034355],
            [1.0, 0.9227477, 1.0, 0.6780561],
            [1.0, 0.91646177, 1.0, 0.6601221],
            [1.0, 0.91059643, 1.0, 0.6468646],
            [1.0, 0.90513545, 1.0, 0.6384181],
            [1.0, 0.9000414, 1.0, 0.63476765],
            [1.0, 0.8952549, 1.0, 0.6357777],
            [1.0, 0.89069635, 1.0, 0.64121217],
            [1.0, 0.8862596, 1.0, 0.6507421],
            [1.0, 0.8818036, 1.0, 0.66393715],
            [1.0, 0.8771375, 1.0, 0.68025357],
            [1.0, 0.8720022, 1.0, 0.6990145],
            [1.0, 0.8660441, 1.0, 0.719399],
            [1.0, 0.85878307, 1.0, 0.7404561],
            [1.0, 0.84957653, 1.0, 0.7611621],
            [1.0, 0.83757895, 1.0, 0.78056324],
            [1.0, 0.8216941, 1.0, 0.7980228],
            [1.0, 0.8005113, 1.0, 0.8135467],
            [0.99108833, 0.77218705, 1.0, 0.828064],
            [0.97618765, 0.7336579, 1.0, 0.8425417],
            [0.9570234, 0.6814261, 1.0, 0.8581086],
            [0.92717046, 0.6107281, 1.0, 0.8740242],
            [0.87264377, 0.51361734, 1.0, 0.8831364],
            [0.7667494, 0.37425363, 1.0, 0.86565703],
            [0.58684665, 0.15105894, 1.0, 0.792903],
            [0.40309954, -0.22192897, 1.0, 0.6778493],
            [0.30307913, -0.6645406, 0.32032692, 0.5561814],
            [0.1432502, -1.0, -0.7834326, 0.30731094],
            [-0.05076139, -1.0, -1.0, -0.18850122],
            [-0.2028995, -1.0, -1.0, -0.47833002],
            [-0.33243275, -1.0, -1.0, -0.63186795],
            [-0.43252927, -1.0, -0.93684345, -0.7109936],
            [-0.50198543, -1.0, -0.8966909, -0.7451998],
            [-0.55477273, -1.0, -0.87718576, -0.7572431],
            [-0.59963316, -1.0, -0.8707306, -0.7596289],
            [-0.641077, -1.0, -0.8736891, -0.7593429],
            [-0.68137753, -1.0, -0.8838504, -0.76042706],
            [-0.72137207, -1.0, -0.8991675, -0.7649133],
            [-0.76085263, -1.0, -0.9172805, -0.77330756],
            [-0.7989401, -1.0, -0.93559915, -0.78509957],
            [-0.8345167, -1.0, -0.95170647, -0.799334],
            [-0.86661166, -1.0, -0.963775, -0.8150858],
            [-0.8946256, -1.0, -0.97078484, -0.8317086],
            [-0.91837806, -1.0, -0.97249895, -0.84885603],
            [-0.9380378, -1.0, -0.96929866, -0.8663847],
            [-0.9540071, -1.0, -0.96197397, -0.8842322],
            [-0.9668097, -1.0, -0.9515317, -0.90233094],
            [-0.97700363, -1.0, -0.9390431, -0.9205734],
            [-0.98512334, -1.0, -0.92553586, -0.93881696],
            [-0.99164504, -1.0, -0.91191846, -0.9569116],
            [-0.99697274, -1.0, -0.8989406, -0.97473806],
            [-1.0, -1.0, -0.8871773, -0.99223274],
            [-1.0, -1.0, -0.87696224, -1.0],
            [-1.0, -1.0, -0.86855894, -1.0],
            [-1.0, -1.0, -0.8622101, -1.0],
            [-1.0, -1.0, -0.85795015, -1.0],
            [-1.0, -1.0, -0.8557587, -1.0],
            [-1.0, -1.0, -0.8555704, -1.0],
            [-1.0, -1.0, -0.8572821, -1.0],
            [-1.0, -1.0, -0.860755, -1.0],
            [-1.0, -1.0, -0.8658133, -1.0],
            [-1.0, -1.0, -0.87224096, -1.0],
            [-1.0, -1.0, -0.87977535, -1.0],
            [-1.0, -1.0, -0.8881058, -1.0],
            [-1.0, -1.0, -0.89687437, -1.0],
            [-1.0, -1.0, -0.9056818, -1.0],
            [-1.0, -1.0, -0.91409653, -1.0],
            [-1.0, -1.0, -0.9216669, -1.0],
            [-1.0, -1.0, -0.9279278, -1.0],
            [-1.0, -1.0, -0.93239695, -1.0],
            [-0.9961943, -1.0, -0.9345514, -1.0],
            [-0.9834586, -1.0, -0.93362767, -1.0],
            [-0.9671113, -1.0, -0.9284645, -1.0],
            [-0.94588476, -1.0, -0.91775376, -1.0],
            [-0.91785616, -1.0, -0.8995549, -1.0],
            [-0.88016766, -1.0, -0.87093955, -1.0],
            [-0.8287441, -1.0, -0.82749766, -1.0],
            [-0.7582472, -1.0, -0.76269644, -1.0],
            [-0.66290134, -1.0, -0.66715723, -1.0],
            [-0.5392508, -1.0, -0.5280629, -1.0],
            [-0.39078623, -1.0, -0.3290425, -1.0],
            [-0.2295668, -1.0, -0.05206226, -1.0],
            [-0.06826158, -1.0, 0.30915332, -1.0],
            [0.08895309, -1.0, 0.7070197, -1.0],
            [0.2400503, -1.0, 1.0, -1.0],
            [0.3742329, -1.0, 1.0, -1.0],
            [0.48094982, -1.0, 1.0, -1.0],
            [0.56609666, -1.0, 1.0, -1.0],
            [0.63677347, -0.88508135, 1.0, -1.0],
            [0.7058708, -0.694147, 1.0, -1.0],
            [0.7992784, -0.5113944, 1.0, -1.0],
            [0.9422653, -0.2919022, 1.0, -1.0],
        ],
        dtype=np.float32,
    )

    for i in range(len(traj)):
        # 14468 is hover rpm
        drone.set_propellels_rpm((1 + 0.05 * traj[i]) * 14468.429183500699)
        scene.step()


if __name__ == "__main__":
    main()
